def view_classify(img, ps):
    ''' Function for viewing an image and it's predicted classes.
    '''
    ps = ps.cpu().data.numpy().squeeze()
    img = img.cpu()
    
    fig, (ax1, ax2) = plt.subplots(figsize=(6,9), ncols=2)
    ax1.imshow(img.resize_(1, 28, 28).numpy().squeeze())
    ax1.axis('off')
    ax2.barh(np.arange(10), ps)
    ax2.set_aspect(0.1)
    ax2.set_yticks(np.arange(10))
    ax2.set_yticklabels(np.arange(10))
    ax2.set_title('Class Probability')
    ax2.set_xlim(0, 1.1)

    plt.tight_layout()
    
%matplotlib inline
def test_prediction(model, data):
    images, labels = next(iter(data))
    img = images[42].view(1, 784)
    
    with torch.no_grad():
        if args.gpu > -1:
            model = model.cuda()
            img = img.cuda()
            logps = model(img)
        else:
            model = model.cpu()
            img = img.cpu()
            logps = model(img)

    # 网络的输出是对数概率，需要对概率取指数
    ps = torch.exp(logps)
    confid, classes = ps.topk(1)
    print("classes is {classes} confidence is {confid:6.4f}".format(classes=classes.item(), confid=confid.item()))
    view_classify(img.view(1, 28, 28), ps)
test_prediction(model, test_loader)
